Skip to content

Add multi-GPU vLLM CAA steering hook#694

Closed
linmou wants to merge 4 commits into
zjunlp:mainfrom
linmou:pr/vllm-caa-multigpu-hook
Closed

Add multi-GPU vLLM CAA steering hook#694
linmou wants to merge 4 commits into
zjunlp:mainfrom
linmou:pr/vllm-caa-multigpu-hook

Conversation

@linmou

@linmou linmou commented Jun 19, 2026

Copy link
Copy Markdown

Summary

This PR adds a lightweight CAA activation-add hook for vLLM inference, including tensor-parallel worker support.

The hook lets EasyEdit apply already-computed CAA vectors to vLLM-loaded models by installing activation additions on selected decoder layers inside each worker process.

What Changed

  • Added steer/vllm_caa_hooks.py

    • install and clear CAA vectors on vLLM model layers
    • worker RPC helpers for tensor-parallel vLLM engines
    • hook call/configuration stats
    • support for common vLLM worker model layouts
  • Added tests/test_vllm_caa_hooks.py

    • focused hook-only unit tests with fake vLLM-style models/workers
    • no dataset dependency
    • no LLM judge dependency
    • no generated artifacts
  • Added examples/vllm_caa_gpu_e2e.py

    • lightweight optional real-GPU smoke test
    • records baseline, steered, and restored outputs
    • records worker install/clear results and hook stats
  • Added docs/vllm_caa_multigpu_hook.md

    • documents runtime API and lightweight validation commands

Validation

pytest -q tests/test_vllm_caa_hooks.py

Result:

9 passed
python -m compileall steer/vllm_caa_hooks.py examples/vllm_caa_gpu_e2e.py tests/test_vllm_caa_hooks.py

Result: passed.

git diff --check main...HEAD

Result: passed.

Optional GPU smoke command:

CUDA_VISIBLE_DEVICES=0,1 \
VLLM_USE_FLASHINFER_SAMPLER=0 \
VLLM_ALLOW_INSECURE_SERIALIZATION=1 \
python examples/vllm_caa_gpu_e2e.py \
  --model /path/to/model \
  --tensor-parallel-size 2 \
  --layer 12 \
  --multiplier 0.0 \
  --vector-value 0.0 \
  --output /tmp/vllm_caa_gpu_e2e.json \
  --monitor-output /tmp/vllm_caa_gpu_e2e.gpu.csv

Notes

The larger steering-effect consistency experiments used during development are not included in this PR. This keeps the runtime contribution focused on the multi-GPU vLLM CAA hook and lightweight validation.

Related Issue

Closes #695

@HaomingX

HaomingX commented Jul 4, 2026

Copy link
Copy Markdown
Collaborator

Hello, thank you for your contribution. To maintain consistency across our repositories, we have reimplemented this in PR #707.

@HaomingX HaomingX closed this Jul 4, 2026
@linmou

linmou commented Jul 5, 2026

Copy link
Copy Markdown
Author

Hello, thank you for your contribution. To maintain consistency across our repositories, we have reimplemented this in PR #707.

Thanks for reviewing #694 and adapting the design into #707. I reviewed the merged implementation, and I understand the direction: #707 generalized the worker-side activation-add hook into EasyEdit2’s apply/inference stage for multiple activation-addition methods, while keeping vector generation in-process.

I would be happy to contribute follow-up improvements on top of #707, especially:

  1. Robustness improvements for the merged vLLM hook path:

    • preserve worker-side exception details instead of returning only False
    • make missing layer installation fail more explicitly
    • clarify the distinction between in-process vLLM and worker-hook/multiprocessing vLLM
    • add lightweight regression tests for steer/vllm_hooks.py
  2. Documentation/config updates:

    • update the EasyEdit2 vLLM tutorial/config examples to show the new multiprocessing options
    • clarify the recommended generate-then-apply workflow
    • document that vector generation remains in-process while apply/inference can use worker-side hooks

I also want to raise one attribution point. Since #694 was the original prototype/design proposal for the worker-side CAA/vLLM hook, and the maintainers said its logic was manually adapted into #707, I think this contribution should be properly recognized if it materially contributed to the core merged method. I am not asking to rewrite history, but I would appreciate a short acknowledgement in #707, #694, or the relevant vLLM documentation.

@HaomingX

HaomingX commented Jul 6, 2026

Copy link
Copy Markdown
Collaborator

Thank you for raising this issue and helping us identify the problem. We appreciate your contribution and will acknowledge your support in the contributor section. You’re also welcome to submit additional code in the future.

For code contributions, our main requirements are that they align with the existing framework, maintain simplicity, and avoid introducing overly AI-specific code.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add multi-GPU vLLM support for CAA steering hooks

2 participants